import uuid
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import JSONResponse
from paddleocr import PaddleOCR
from PIL import Image

# 尝试导入 paddle（支持 NPU 版本）
try:
    import paddle
    PADDLE_AVAILABLE = True
except ImportError as e:
    PADDLE_AVAILABLE = False
    print(f"警告: paddle 模块导入失败: {e}")
    print("请检查 paddle 是否正确安装:")
    print("  - 对于 NPU 版本: pip install paddle-custom-npu -i https://www.paddlepaddle.org.cn/packages/nightly/npu")
    print("  - 对于 CPU/GPU 版本: pip install paddlepaddle")
    paddle = None
try:
    import pymupdf as fitz  # PyMuPDF - 使用 pymupdf 包名导入
except ImportError:
    # 如果 pymupdf 不可用，尝试直接导入 fitz（旧版本的 PyMuPDF）
    try:
        import fitz
    except ImportError:
        raise ImportError(
            "PyMuPDF 未安装。请运行: pip install pymupdf 或 pip install PyMuPDF"
        )

# 初始化 FastAPI 应用
app = FastAPI(title="OCR识别API (PP-OCRv5)", description="支持文件上传进行OCR识别，使用PP-OCRv5版本")

# 全局初始化 OCR 模型（避免重复加载）
ocr_instance = None


def init_ocr_model():
    """初始化 OCR 模型"""
    global ocr_instance
    if ocr_instance is None:
        # 检查设备信息并设置 NPU
        print("=" * 80)
        print("设备信息检查:")
        print("=" * 80)
        
        if not PADDLE_AVAILABLE or paddle is None:
            print("警告: paddle 模块不可用，无法检查设备信息")
            print("请确保已正确安装 paddle:")
            print("  - NPU 版本: pip install paddle-custom-npu -i https://www.paddlepaddle.org.cn/packages/nightly/npu")
            print("  - CPU/GPU 版本: pip install paddlepaddle")
        else:
            try:
                # 设置使用 NPU 设备（优先使用 npu:0）
                try:
                    paddle.set_device('npu:0')
                    print("✓ 已设置设备为 npu:0")
                except Exception as e:
                    print(f"⚠ 设置 npu:0 失败，尝试使用 npu: {e}")
                    try:
                        paddle.set_device('npu')
                        print("✓ 已设置设备为 npu")
                    except Exception as e2:
                        print(f"⚠ 设置 npu 失败: {e2}")
                        print("使用默认设备")
                
                # 获取当前设备
                current_device = paddle.device.get_device()
                print(f"Current device: {current_device}")
                
                # 检查是否可用NPU（某些版本可能没有这个方法）
                try:
                    if hasattr(paddle.device, 'is_compiled_with_npu'):
                        is_npu_available = paddle.device.is_compiled_with_npu()
                        print(f"NPU compiled: {is_npu_available}")
                    else:
                        # 通过尝试设置 NPU 来判断
                        if 'npu' in str(current_device).lower():
                            print("NPU compiled: True (从设备名称推断)")
                        else:
                            print("NPU compiled: Unknown (无法检测)")
                except Exception as e:
                    print(f"NPU available: 无法检测 (方法不存在: {type(e).__name__})")
                
                # 检查是否可用GPU
                try:
                    if hasattr(paddle.device, 'is_compiled_with_cuda'):
                        is_gpu_available = paddle.device.is_compiled_with_cuda()
                        print(f"GPU compiled: {is_gpu_available}")
                    else:
                        print("GPU compiled: Unknown (无法检测)")
                except Exception as e:
                    print(f"GPU available: 无法检测 ({e})")
                
                print("=" * 80)
            except Exception as e:
                print(f"设备检查失败: {e}")
                import traceback
                traceback.print_exc()
        
        # 初始化 OCR 模型（会自动使用 paddle.set_device() 设置的设备，即 npu:0）
        # 注意：PaddleOCR 会自动使用 paddle 设置的设备，无需额外参数
        ocr_instance = PaddleOCR(
            use_doc_orientation_classify=False,  # 禁用文档方向分类模块
            use_doc_unwarping=False,  # 禁用文本图像矫正模块
            use_textline_orientation=False,  # 禁用文本行方向分类模块
            ocr_version="PP-OCRv5"  # 明确指定使用 PP-OCRv5 版本
        )
    return ocr_instance


def pdf_to_images(pdf_path, output_dir, dpi=200):
    """
    将PDF转换为图片列表

    Args:
        pdf_path: PDF文件路径
        output_dir: 输出目录
        dpi: 转换分辨率（默认200）

    Returns:
        list: 图片路径列表
    """
    image_paths = []
    try:
        pdf_document = fitz.open(pdf_path)
        for page_num in range(len(pdf_document)):
            page = pdf_document[page_num]
            # 设置缩放比例（DPI）
            zoom = dpi / 72.0
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)

            # 保存为PNG图片
            image_path = output_dir / f"page_{page_num + 1}.png"
            pix.save(image_path)
            image_paths.append(image_path)

        pdf_document.close()
        return image_paths
    except Exception as e:
        print(f"PDF转换失败: {e}")
        raise HTTPException(status_code=500, detail=f"PDF转换失败: {str(e)}")


def preprocess_image(image_path, max_size=(4000, 4000), max_megapixels=16):
    """
    预处理图片：如果图片太大则压缩，避免内存溢出和处理超时

    Args:
        image_path: 图片路径
        max_size: 最大尺寸 (width, height)
        max_megapixels: 最大百万像素数（16MP = 约4000x4000）

    Returns:
        tuple: (处理后的图片路径, 是否压缩, 原始文件大小(字节), 压缩后文件大小(字节))
    """
    try:
        original_size = image_path.stat().st_size if image_path.exists() else 0
        
        with Image.open(image_path) as img:
            width, height = img.size
            megapixels = (width * height) / 1_000_000

            # 如果图片太大，进行压缩
            if megapixels > max_megapixels or width > max_size[0] or height > max_size[1]:
                # 计算缩放比例
                scale = min(max_size[0] / width, max_size[1] / height,
                            (max_megapixels * 1_000_000 / (width * height)) ** 0.5)
                new_width = int(width * scale)
                new_height = int(height * scale)

                # 压缩图片
                img_resized = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

                # 保存压缩后的图片
                processed_path = image_path.parent / f"{image_path.stem}_processed{image_path.suffix}"
                img_resized.save(processed_path, quality=95, optimize=True)
                
                compressed_size = processed_path.stat().st_size if processed_path.exists() else 0

                return processed_path, True, original_size, compressed_size  # 返回新路径、是否压缩标志、原始大小、压缩后大小
            else:
                return image_path, False, original_size, original_size  # 原路径，未压缩，大小相同
    except Exception as e:
        print(f"图片预处理失败，使用原图: {e}")
        original_size = image_path.stat().st_size if image_path.exists() else 0
        return image_path, False, original_size, original_size


@app.get("/")
async def root():
    """根路径，返回 API 信息"""
    return {"message": "OCR识别API (PP-OCRv5)", "status": "running", "version": "PP-OCRv5"}


@app.post("/ocr")
async def upload_and_ocr(
    file: UploadFile = File(...),
    enable_compression: str = Query("false", description="是否启用图片压缩，默认为false")
):
    """
    上传文件并进行OCR识别（支持图片和PDF扫描件）

    Args:
        file: 上传的文件（支持图片和PDF）
        enable_compression: 是否启用图片压缩，默认为false（字符串）

    Returns:
        JSONResponse: 包含OCR识别结果的响应
    """
    # 使用当前目录下的 tmp 文件夹
    base_dir = Path(__file__).parent.parent  # 获取项目根目录
    tmp_dir = base_dir / "tmp"
    tmp_dir.mkdir(exist_ok=True)

    # 判断是否启用压缩（将字符串转换为布尔值）
    should_compress = enable_compression.lower() == "true"

    # 生成唯一文件名
    file_id = str(uuid.uuid4())
    file_ext = Path(file.filename).suffix.lower() if file.filename else ''
    input_file_path = tmp_dir / f"{file_id}{file_ext}"
    processed_file_paths = []  # 存储所有处理后的文件路径（包括PDF转换的图片）
    was_compressed = False
    original_image_size = 0  # 原始图片大小（字节）
    compressed_image_size = 0  # 压缩后图片大小（字节）
    is_pdf = False
    pdf_pages = 0

    try:
        # 保存上传的文件
        with open(input_file_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)

        # 初始化 OCR 模型
        ocr = init_ocr_model()

        # 判断文件类型并处理
        if file_ext == '.pdf':
            # PDF文件：转换为图片后逐页OCR
            is_pdf = True
            pdf_images_dir = tmp_dir / f"{file_id}_pages"
            pdf_images_dir.mkdir(exist_ok=True)

            # 将PDF转换为图片
            image_paths = pdf_to_images(input_file_path, pdf_images_dir)
            pdf_pages = len(image_paths)

            # 对所有页面进行OCR识别
            all_results = []
            all_texts = []
            page_results = []

            for page_idx, image_path in enumerate(image_paths):
                try:
                    # 如果启用压缩，预处理图片（压缩）
                    if should_compress:
                        processed_path, compressed, orig_size, comp_size = preprocess_image(image_path)
                        if compressed:
                            processed_file_paths.append(processed_path)
                            # 记录压缩信息（使用第一个压缩的图片的信息，或累计）
                            if not was_compressed:
                                original_image_size = orig_size
                                compressed_image_size = comp_size
                                was_compressed = True
                        ocr_input_path = str(processed_path) if compressed else str(image_path)
                    else:
                        processed_path = image_path
                        compressed = False
                        ocr_input_path = str(image_path)
                    
                    # 始终记录原始图片路径以便删除
                    processed_file_paths.append(image_path)

                    # 执行OCR识别
                    result = ocr.predict(ocr_input_path)

                    # 格式化该页的结果
                    page_texts = []
                    for res in result:
                        rec_texts = res.get("rec_texts", []) if isinstance(res, dict) else []
                        page_texts.extend(rec_texts)
                        all_texts.extend(rec_texts)

                        # 收集所有结果用于最终格式化
                        all_results.append(res)

                    # 记录每页的结果
                    page_results.append({
                        "page": page_idx + 1,
                        "texts": page_texts,
                        "full_text": "\n".join(page_texts),
                        "text_count": len(page_texts)
                    })
                except Exception as e:
                    print(f"第 {page_idx + 1} 页OCR失败: {e}")
                    page_results.append({
                        "page": page_idx + 1,
                        "error": str(e),
                        "texts": [],
                        "full_text": "",
                        "text_count": 0
                    })

            # 格式化最终结果
            ocr_results = []
            for res in all_results:
                rec_texts = res.get("rec_texts", []) if isinstance(res, dict) else []
                result_item = {
                    "rec_texts": rec_texts,
                    "full_text": "\n".join(rec_texts) if rec_texts else ""
                }
                if isinstance(res, dict):
                    if "dt_boxes" in res:
                        result_item["dt_boxes"] = res["dt_boxes"]
                    if "rec_scores" in res:
                        result_item["rec_scores"] = res["rec_scores"]
                    # 添加所有原始字段（除了已处理的）
                    for key, value in res.items():
                        if key not in ["rec_texts", "dt_boxes", "rec_scores"]:
                            result_item[key] = value
                else:
                    result_item["raw_result"] = str(res)
                ocr_results.append(result_item)

            full_text = "\n".join(all_texts)

        else:
            # 图片文件处理
            # 如果是图片文件且启用压缩，进行预处理（压缩大图片）
            if should_compress and file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp']:
                processed_file_path, was_compressed, original_image_size, compressed_image_size = preprocess_image(input_file_path)
                if was_compressed:
                    processed_file_paths.append(processed_file_path)
                ocr_input_path = str(processed_file_path) if was_compressed else str(input_file_path)
            else:
                ocr_input_path = str(input_file_path)

            # 执行 OCR 识别（增加超时保护）
            try:
                import signal
                import platform

                # 只在非Windows系统上使用signal超时
                if platform.system() != 'Windows':
                    def timeout_handler(signum, frame):
                        raise TimeoutError("OCR处理超时")

                    signal.signal(signal.SIGALRM, timeout_handler)
                    signal.alarm(300)  # 5分钟超时
                    result = ocr.predict(ocr_input_path)
                    signal.alarm(0)  # 取消超时
                else:
                    # Windows系统直接执行（可通过uvicorn timeout设置超时）
                    result = ocr.predict(ocr_input_path)
            except TimeoutError:
                raise HTTPException(status_code=504, detail="OCR处理超时，图片可能过大或文本过长，建议压缩图片后重试")
            except Exception as ocr_error:
                error_msg = str(ocr_error)
                # 如果是内存相关错误，给出更友好的提示
                if "memory" in error_msg.lower() or "内存" in error_msg or "OOM" in error_msg:
                    raise HTTPException(
                        status_code=507,
                        detail="内存不足，图片太大或文本过长。已自动压缩图片，如仍有问题请尝试更小的图片。"
                    )
                raise HTTPException(status_code=500, detail=f"OCR识别失败: {error_msg}")

            # 格式化OCR结果
            ocr_results = []
            all_texts = []

            for res in result:
                # 提取识别文本
                rec_texts = res.get("rec_texts", []) if isinstance(res, dict) else []

                # 构建结果结构
                result_item = {
                    "rec_texts": rec_texts,  # 识别文本列表
                    "full_text": "\n".join(rec_texts) if rec_texts else ""  # 完整文本（换行连接）
                }

                # 如果有其他信息（坐标、置信度等），也添加进去
                if isinstance(res, dict):
                    if "dt_boxes" in res:
                        result_item["dt_boxes"] = res["dt_boxes"]
                    if "rec_scores" in res:
                        result_item["rec_scores"] = res["rec_scores"]
                    # 添加所有原始字段（除了已处理的）
                    for key, value in res.items():
                        if key not in ["rec_texts", "dt_boxes", "rec_scores"]:
                            result_item[key] = value
                else:
                    # 如果不是字典，尝试获取所有属性
                    result_item["raw_result"] = str(res)

                ocr_results.append(result_item)
                all_texts.extend(rec_texts)

            full_text = "\n".join(all_texts)
            page_results = None  # 图片文件没有分页结果

        # 检查结果大小，如果文本过长可能存在问题
        if len(full_text) > 100000:  # 如果文本超过10万字符，给出警告
            print(f"警告: 识别文本过长 ({len(full_text)} 字符)")

        # 打印识别的结果
        print("=" * 80)
        print("OCR识别结果 (full_text):")
        print("=" * 80)
        print(full_text)
        print("=" * 80)
        
        # 如果压缩了图片，打印压缩前后的大小
        if was_compressed:
            def format_size(size_bytes):
                """格式化文件大小"""
                for unit in ['B', 'KB', 'MB', 'GB']:
                    if size_bytes < 1024.0:
                        return f"{size_bytes:.2f} {unit}"
                    size_bytes /= 1024.0
                return f"{size_bytes:.2f} TB"
            
            print(f"图片已压缩:")
            print(f"  压缩前大小: {format_size(original_image_size)} ({original_image_size} 字节)")
            print(f"  压缩后大小: {format_size(compressed_image_size)} ({compressed_image_size} 字节)")
            compression_ratio = (1 - compressed_image_size / original_image_size) * 100 if original_image_size > 0 else 0
            print(f"  压缩率: {compression_ratio:.2f}%")

        # 返回结果
        response_data = {
            "status": "success",
            "file_id": file_id,
            "original_filename": file.filename,
            "file_type": "pdf" if is_pdf else "image",
            "results": ocr_results,
            "all_texts": all_texts,  # 所有识别文本的列表
            "full_text": full_text,  # 完整文本（所有文本用换行连接）
            "text_count": len(all_texts),
            "result_count": len(ocr_results),
            "image_compressed": was_compressed,  # 是否压缩了图片
            "text_length": len(full_text)  # 文本总长度
        }

        # 如果压缩了图片，添加压缩信息到响应中
        if was_compressed:
            response_data["original_image_size"] = original_image_size  # 原始大小（字节）
            response_data["compressed_image_size"] = compressed_image_size  # 压缩后大小（字节）
            response_data["compression_ratio"] = round((1 - compressed_image_size / original_image_size) * 100, 2) if original_image_size > 0 else 0  # 压缩率

        # PDF特有信息
        if is_pdf:
            response_data["pdf_pages"] = pdf_pages
            response_data["page_results"] = page_results

        return JSONResponse(response_data["full_text"])

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"OCR处理失败: {str(e)}")

    finally:
        # 删除上传的文件和处理后的文件（包括PDF转换的图片）
        try:
            if input_file_path.exists():
                input_file_path.unlink()

            # 删除所有处理后的文件
            for file_path in processed_file_paths:
                if file_path and file_path.exists():
                    file_path.unlink()

            # 删除PDF页面图片目录
            if is_pdf:
                pdf_images_dir = tmp_dir / f"{file_id}_pages"
                if pdf_images_dir.exists():
                    import shutil
                    shutil.rmtree(pdf_images_dir, ignore_errors=True)
        except Exception as e:
            print(f"删除文件失败: {e}")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8869)

